# Copyright (c) HySoP 2011-2024
#
# This file is part of HySoP software.
# See "https://particle_methods.gricad-pages.univ-grenoble-alpes.fr/hysop-doc/"
# for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
@file transpose.py
Transposition related operators.
"""
from hysop.constants import Implementation, Backend, DirectionLabels
from hysop.tools.htypes import check_instance, to_tuple
from hysop.tools.decorators import debug
from hysop.tools.transposition_states import TranspositionState
from hysop.fields.continuous_field import Field, ScalarField
from hysop.topology.cartesian_descriptor import CartesianTopologyDescriptors
from hysop.core.graph.node_generator import ComputationalGraphNodeGenerator
from hysop.core.graph.computational_operator import ComputationalGraphNode
from hysop.testsenv import __HAS_OPENCL_BACKEND__
[docs]
class TranspositionNotImplementedError(NotImplementedError):
"""
Error raised for unimplemented transposition operator configurations.
"""
pass
[docs]
class Transpose(ComputationalGraphNodeGenerator):
"""
Operator generator for inplace and out of place field transposition
and permutations in general.
Available implementations are:
*python (numpy based transpose, n-dimensional)
*opencl (opencl codegen based transpose up to 16D)
Implementations handle only one field at a time but this
graph node generator generates an operator per supplied field,
possibly on different implementation backends.
See hysop.operator.base.transpose_operator.TransposeOperatorBase for operator backend
implementation interface.
"""
[docs]
@classmethod
def implementations(cls):
from hysop.backend.host.python.operator.transpose import PythonTranspose
_implementations = {
Implementation.PYTHON: PythonTranspose,
}
if __HAS_OPENCL_BACKEND__:
from hysop.backend.device.opencl.operator.transpose import OpenClTranspose
_implementations.update(
{
Implementation.OPENCL: OpenClTranspose,
}
)
return _implementations
[docs]
@classmethod
def default_implementation(cls):
msg = "Transpose has no default implementation, "
msg += "implementation should match the discrete field topology backend."
raise RuntimeError(msg)
@debug
def __new__(
cls,
fields,
variables,
axes,
output_fields=None,
implementation=None,
name=None,
base_kwds=None,
**kwds,
):
base_kwds = {} if (base_kwds is None) else {}
return super().__new__(
cls,
name=name,
candidate_input_tensors=None,
candidate_output_tensors=None,
**base_kwds,
)
@debug
def __init__(
self,
fields,
variables,
axes,
output_fields=None,
implementation=None,
name=None,
base_kwds=None,
**kwds,
):
"""
Initialize a Transpose operator generator operating on CartesianTopology topologies.
Parameters
----------
fields: Field, list or tuple of Fields
Input continuous fields to be transposed, at least 2D.
All fields should have the same dimension.
output_fields: Field, list or tuple of Fields, optional
Output continuous fields where the results are stored.
Transposed shapes should match the input.
By default output_fields are the same as input_fields
resulting in inplace transpositions.
Input and output are matched by order int list/tuple.
variables: dict
Dictionary of fields as keys and CartesianTopologyDescriptors as values.
axes: tuple of ints, or array like of tuples, or dict of (tuple, TranspositionState).
Permutation of axes in numpy notations (as a tuple of ints).
Axe dim-1 is the contiguous axe, axe 0 has the greatest stride in memory.
It can also be an array like of axes and the operator implementation will choose
the most suitable axe permutation scheme in this list.
If a dictionnary is provided, this gives the operator implementation to choose
the most suitable axe permutation scheme with the knowledge of the target
transposition state as well.
All fields on the same backend will perform the same permutation.
Operator chosen permutation can be retrieved by using generated
operator 'axes' attribute.
implementation: Implementation, optional, defaults to None
target implementation, should be contained in available_implementations().
If implementation is set and topology does not match backend,
RuntimeError will be raised on _generate.
If None, implementation will be set according to topologies backend,
different implementations may be choosen for different Fields if
defined on different backends.
name: string
prefix for generated operator names
base_kwds: dict, optional, defaults to None
Base class keywords arguments.
If None, an empty dict will be passed.
kwds:
keywords arguments that will be passed towards implementation
transpose operator __init__.
Notes
-----
Out of place transpose will always be faster to process.
The only exception to this rule may be 2D square matrices.
Inplace transposition may request a temporary buffer because not all implementations
may support inplace transposition.
* Example for axes:
- In 4D, axes=(0,1,2,3) means no permutations at all, this will raise ValueError.
- In 2D, axes=(1,0) is a transposition.
- In 3D, axes=(0,2,1) is a transposition in the XY-plane.
* About dimensions:
- No limit for the numpy CPU transposition implementation.
- Max 16D for the OpenCL implementation.
A Transpose operator implementation should support the TransposeOperatorBase
interface (see hysop.operator.base.transpose_operator.TransposeOperatorBase).
This ComputationalGraphNodeFrontend will generate a operator for each
input and output ScalarField pair.
All implementations should raise TranspositionNotImplementedError is the user supplied
parameters leads to unimplemented or unsupported transposition features.
"""
input_fields = to_tuple(fields)
assert None not in input_fields
if output_fields is not None:
output_fields = to_tuple(output_fields)
output_fields = tuple(
ofield if (ofield is not None) else ifield
for (ifield, ofield) in zip(input_fields, output_fields)
)
else:
output_fields = tuple(ifield for ifield in input_fields)
check_instance(input_fields, tuple, values=Field)
check_instance(output_fields, tuple, values=Field, size=len(input_fields))
candidate_input_tensors = tuple(filter(lambda x: x.is_tensor, input_fields))
candidate_output_tensors = tuple(filter(lambda x: x.is_tensor, output_fields))
base_kwds = {} if (base_kwds is None) else {}
super().__init__(
name=name,
candidate_input_tensors=candidate_input_tensors,
candidate_output_tensors=candidate_output_tensors,
**base_kwds,
)
# expand tensors
ifields, ofields = (), ()
for ifield, ofield in zip(input_fields, output_fields):
msg = "Input and output field shape mismatch, got field {} of shape {} "
msg += "and field {} of shape {}."
if ifield.is_tensor ^ ofield.is_tensor:
if ifield.is_tensor:
msg = msg.format(
ifield.short_description(),
ifield.shape,
ofield.short_description(),
"(1,)",
)
else:
msg = msg.format(
ifield.short_description(),
"(1,)",
ofield.short_description(),
ofield.shape,
)
raise RuntimeError(msg)
if ifield.is_tensor and ofield.is_tensor and (ifield.shape != ofield.shape):
msg = msg.format(
ifield.short_description(),
ifield.shape,
ofield.short_description(),
ofield.shape,
)
raise RuntimeError(msg)
ifields += ifield.fields
ofields += ofield.fields
input_fields = ifields
output_fields = ofields
check_instance(input_fields, tuple, values=ScalarField)
check_instance(output_fields, tuple, values=ScalarField, size=len(input_fields))
check_instance(variables, dict, keys=Field, values=CartesianTopologyDescriptors)
check_instance(base_kwds, dict, keys=str)
check_instance(name, str, allow_none=True)
fields = set(input_fields).union(output_fields)
vfields = {f for tfield in variables.keys() for f in tfield.fields}
if not fields:
raise ValueError("fields are empty.")
if not vfields:
raise ValueError("variables are empty.")
if not axes:
raise ValueError("axes are empty.")
if fields != vfields:
if fields.difference(vfields):
missing = tuple(field.name for field in fields - vfields)
msg = "Missing fields in variables parameter: {}"
msg = msg.format(", ".join(missing))
else:
missing = tuple(field.name for field in vfields - fields)
msg = "Too many fields present in variables keys: {}"
msg = msg.format(", ".join(missing))
raise ValueError(msg)
dim = input_fields[0].domain.dim
if dim < 2:
msg = "input_field dimension should be at least 2 but {} is a {}D field."
msg = msg.format(input_fields[0].name, dim)
raise ValueError(msg)
for input_field, output_field in zip(input_fields, output_fields):
in_topo_descriptor = ComputationalGraphNode.get_topo_descriptor(
variables, input_field
)
out_topo_descriptor = ComputationalGraphNode.get_topo_descriptor(
variables, output_field
)
if input_field.domain != output_field.domain:
msg = "input_field {} and output_field {} do not share the same domain."
msg.format(input_field.name, output_field.name)
raise ValueError(msg)
idim = input_field.domain.dim
if idim != dim:
msg = "input_field {} is of dimension {} and does not match first "
msg += "input_field {} dimension {}."
msg.format(input_field, idim, input_fields[0].name, dim)
raise ValueError(msg)
candidate_axes = self._check_axes(axes, dim)
self.input_fields = input_fields
self.output_fields = output_fields
self.variables = variables
self.candidate_axes = candidate_axes
self.implementation = implementation
self.kwds = kwds
@staticmethod
def _check_axes(axes, dim):
msg = "Unkown type for axes parameter."
if isinstance(axes, dict):
candidate_axes = axes
elif isinstance(axes, (list, tuple, set, frozenset)):
if isinstance(axes[0], int):
msg = "Axes should be ordered but got a set."
assert not isinstance(axes, (set, frozenset)), msg
candidate_axes = (tuple(axes),)
elif isinstance(axes[0], (list, tuple)):
candidate_axes = tuple(tuple(ax) for ax in axes)
else:
raise TypeError(msg)
candidate_axes = dict(zip(candidate_axes, (None,) * len(candidate_axes)))
else:
raise TypeError(msg)
del axes
check_instance(candidate_axes, dict, keys=tuple)
for axes, target_tstate in candidate_axes.items():
check_instance(axes, tuple, values=int)
check_instance(target_tstate, TranspositionState[dim], allow_none=True)
if len(axes) != dim:
msg = "All axes should have the dimension of the transposed fields {} "
msg += "but got axes={} which is of dimension {}."
msg = msg.format(dim, axes, len(axes))
raise ValueError(msg)
if set(axes) != set(range(dim)):
msg = f"Axes should describe a valid pemutation, got axes={axes}."
raise ValueError(msg)
if axes == tuple(range(dim)):
msg = "Given axes is identity (no permutation)."
raise ValueError(msg)
return candidate_axes
def _get_op_and_check_implementation(self, src_topo, dst_topo):
backend = getattr(src_topo, "backend", None) or getattr(
dst_topo, "backend", None
)
if backend is None:
implementation = self.implementation
if implementation is None:
msg = (
"Source and destination topology descriptors do not expose backend "
)
msg += (
"attribute and self.implementation was not set, cannot determine "
)
msg += "the target transposition operator."
raise ValueError(msg)
if implementation not in self.implementations():
msg = f"Unknown transposition implementation {implementation}."
raise ValueError(msg)
op_cls = self.implementations()[implementation]
else:
def check_impl(impl, allowed):
if impl not in allowed:
msg = "Specified transposition implementation {} but this is not a valid "
msg += "implementation for fields defined on backend of kind {}."
msg = msg.format(impl, backend.kind)
raise ValueError(msg)
if backend.kind == Backend.OPENCL:
implementation = self.implementation or Implementation.OPENCL
check_impl(implementation, (Implementation.OPENCL,))
op_cls = self.implementations()[implementation]
elif backend.kind == Backend.HOST:
implementation = self.implementation or Implementation.PYTHON
check_impl(implementation, (Implementation.PYTHON,))
op_cls = self.implementations()[implementation]
else:
assert implementation not in self.implementations()
msg = f"Unknown transposition implementation {implementation}."
raise ValueError(msg)
from hysop.operator.base.transpose_operator import TransposeOperatorBase
if not issubclass(op_cls, TransposeOperatorBase):
msg = "Class {} does not inherit from the transpose operator interface "
msg += "({}). This is an implementation error."
msg = msg.format(op_cls, TransposeOperatorBase)
raise TypeError(msg)
if not issubclass(op_cls, ComputationalGraphNode):
msg = (
"Class {} does not inherit from the computational graph node interface"
)
msg += "({}). This is an implementation error."
msg = msg.format(op_cls, ComputationalGraphNode)
raise TypeError(msg)
return op_cls
@debug
def _generate(self):
nodes = []
for ifield, ofield in zip(self.input_fields, self.output_fields):
src_topo = ComputationalGraphNode.get_topo_descriptor(
self.variables, ifield
)
dst_topo = ComputationalGraphNode.get_topo_descriptor(
self.variables, ofield
)
kwds = self.kwds.copy()
TransposeOp = self._get_op_and_check_implementation(src_topo, dst_topo)
axes = TransposeOp.get_preferred_axes(
src_topo, dst_topo, self.candidate_axes
)
variables = {ifield: src_topo}
variables[ofield] = dst_topo
# instantiate operator
node = TransposeOp(
input_field=ifield,
output_field=ofield,
variables=variables,
axes=axes,
**kwds,
)
nodes.append(node)
return nodes